package com.ficus.road.maintain.util;

import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.BiFunction;

@Slf4j
public class BatchUtil {

    public static <K, T, R>
    BiFunction<K, T, CompletableFuture<R>> createBatchFunc(
            BiFunction<K, List<T>, List<R>> batchFunc,
            BiFunction<K, List<T>, Boolean> batchChecker,
            int timeoutInMs,
            Executor executor) {
        return (new InternalBatchExecutor<>(
                timeoutInMs, batchFunc, batchChecker, executor));
    }

    private static class BatchEntry<T, R> {
        ArrayList<T> items;
        ArrayList<CompletableFuture<R>> futures;
        ReentrantLock lock;
        // 这个字段初始化为false 一旦batch满了或者超时了就置为true
        // 一旦置为true就不会再改变了, 表示之后马上要进入batchExecutor运行
        boolean needToRun;

        BatchEntry() {
            this.items = new ArrayList<>();
            this.futures = new ArrayList<>();
            this.lock = new ReentrantLock();
            this.needToRun = false;
        }
    }

    private static class InternalBatchExecutor<K, T, R>
            implements BiFunction<K, T, CompletableFuture<R>>{
        private TimeoutMap<K, BatchEntry<T, R>> timeoutMap;
        private BiFunction<K, List<T>, List<R>> batchFunc;
        private BiFunction<K, List<T>, Boolean> batchChecker;
        private Executor batchExecutor;

        InternalBatchExecutor(
                int timeoutInMs,
                BiFunction<K, List<T>, List<R>> batchFunc,
                BiFunction<K, List<T>, Boolean> batchChecker,
                Executor batchExecutor) {
            this.timeoutMap = new TimeoutMap<>(timeoutInMs, this::timeoutFunc);
            this.batchFunc = batchFunc;
            this.batchChecker = batchChecker;
            this.batchExecutor = batchExecutor;
        }

        @Override
        public CompletableFuture<R> apply(K key, T item) {
            CompletableFuture<R> future = null;
            while (true) {
                future = internalApply(key, item);
                if (future == null) {
                    Thread.yield();
                } else {
                    return future;
                }
            }
        }

        private CompletableFuture<R> internalApply(K key, T item) {

            BatchEntry<T, R> entry = timeoutMap.get(key);
            if (entry == null) {
                entry = timeoutMap.getOrDefault(key, new BatchEntry<>());
            }
            boolean needToRun = false;
            entry.lock.lock();
            CompletableFuture<R> future = null;
            try {
                // 刚好其他线程 将batch凑满 准备马上删除 需要去重试
                if (entry.needToRun) {
                    return null;
                }
                entry.items.add(item);
                boolean isFull = false;
                try {
                    Boolean tmpFull = batchChecker.apply(key, entry.items);
                    if (tmpFull != null && tmpFull) {
                        isFull = true;
                    }
                } catch (Exception e) {
                    log.error("batch checker error!", e);
                    isFull = true;
                }
                if (isFull) {
                    // 将刚刚插入的最后一个弹出来 放到下一个新的batch里面
                    entry.items.remove(entry.items.size() - 1);
                    // 如果batch凑满了，就标记为 需要运行状态
                    entry.needToRun = true;
                    needToRun = true;
                } else {
                    future = new CompletableFuture<>();
                    entry.futures.add(future);
                }
            } finally {
                entry.lock.unlock();
            }
            if (needToRun) {
                // 先删除从而可以其他线程在这个key上继续攒batch
                timeoutMap.remove(key);
                runEntry(key, entry);
            }
            return future;
        }

        private void timeoutFunc(K key, BatchEntry<T, R> entry) {
            boolean needToRun = false;
            entry.lock.lock();
            try {
                if (!(entry.needToRun)) {
                    entry.needToRun = true;
                    needToRun = true;
                }
            } finally {
                entry.lock.unlock();
            }
            if (needToRun) {
                runEntry(key, entry);
            }
        }

        private void runEntry(K key, BatchEntry<T, R> entry) {
            final List<T> items = entry.items;
            final List<CompletableFuture<R>> futures = entry.futures;
            batchExecutor.execute(() -> {
                List<R> results = null;
                try {
                    results = batchFunc.apply(key, items);
                } catch (Exception e) {
                    log.error("run batch function error!", e);
                }
                if (results == null || results.size() != items.size()) {
                    Iterator<CompletableFuture<R>> futureIterator =
                            futures.iterator();
                    while (futureIterator.hasNext()) {
                        futureIterator.next().completeExceptionally(
                                new RuntimeException("batch function error"));
                    }
                } else {
                    Iterator<CompletableFuture<R>> futureIterator =
                            futures.iterator();
                    Iterator<R> resultIterator= results.iterator();
                    while (futureIterator.hasNext()) {
                        futureIterator.next().complete(resultIterator.next());
                    }
                }
            });
        }
    }
}
